package ai.tuobot.sdk.util

import android.content.Context
import android.util.Log
import kotlinx.coroutines.CancellableContinuation
import kotlinx.coroutines.suspendCancellableCoroutine
import org.json.JSONException
import org.json.JSONObject
import org.vosk.Model
import org.vosk.Recognizer
import org.vosk.android.RecognitionListener
import org.vosk.android.SpeechStreamService
import org.vosk.android.StorageService
import java.io.ByteArrayInputStream
import java.io.IOException
import java.io.InputStream
import kotlin.coroutines.cancellation.CancellationException
import kotlin.coroutines.resume

/**
 * 本地语音转文字
 */
class LocalSTT(val context: Context){
    private val SAMPLE_RATE = 16000f

    companion object{
        val TAG = "LocalSTT"
        val STATE_NOT_READY = -1
        val STATE_READY = 0
        val STATE_ERROR_UNPACK = 1
    }

    private lateinit var cont: CancellableContinuation<String>
    private var model: Model? = null

    private var currentState = STATE_NOT_READY

    private var recognizer: Recognizer? = null

    suspend fun initModel():Int {
        return suspendCancellableCoroutine { cont ->
            StorageService.unpack(context, "model-en-us", "model",  //vosk-model-small-en-us-0.15
                {
                    model: Model -> this.model = model
                    currentState = STATE_READY
                    cont.resume(currentState)
                },
                { exception: IOException ->
                    currentState =  STATE_ERROR_UNPACK
                    cont.resume(currentState)
                }
            )
        }
    }

    /**
     * Preheating model
     */
    suspend fun preHot() {
        val byteArray = ByteArray(1024*1000) { it.toByte() } // Create a byte array with 50 bytes
        val inputStream = ByteArrayInputStream(byteArray)
        recognizeFile(inputStream)
    }

    suspend fun recognizeFile(ais: InputStream): HashMap<String, ArrayList<JSONObject>> {
        val recMapRes = HashMap<String, ArrayList<JSONObject>>()
        return suspendCancellableCoroutine { cont ->
            try {
                if (ais.skip(44) != 44L){
                    Log.d(TAG, "Error in skip(44) input stream")
                    cont.resume(recMapRes)
                    return@suspendCancellableCoroutine
                }

                model.let { itOfModel ->
                    if (recognizer == null) {
                        recognizer = Recognizer(
                            itOfModel,
                            SAMPLE_RATE
                        ) // "[\"one zero zero zero one\", " + "\"oh zero one two three four five six seven eight nine\", \"[unk]\"]"
                    }
                    val speechStreamService = SpeechStreamService(recognizer, ais, SAMPLE_RATE)
                    //Log.d(TAG, "Start speechStreamService")
                    speechStreamService.start(object : RecognitionListener {

                        private fun prepareWords(hypothesis: String, type:String) {
                            try {
                                val Jbj = JSONObject(hypothesis)
                                var wordsList = recMapRes[type]
                                if (null == wordsList) {
                                    wordsList = ArrayList()
                                    wordsList.add(Jbj)
                                } else {
                                    wordsList.add(Jbj)
                                }
                                recMapRes.put(type, wordsList)
                            } catch (e: JSONException) {
                                e.message?.let { Log.d(TAG, it) }
                            }

                        }

                        override fun onPartialResult(hypothesis: String?) {
                            try {
                                val Jbj = hypothesis?.let { JSONObject(it) }
                                if (Jbj != null) {
                                    if (Jbj.getString("partial").isNotEmpty()) {
                                        prepareWords(hypothesis,"onPartialResult")
                                        speechStreamService.stop()
                                    }
                                }
                            } catch (e: JSONException) {
                                e.message?.let { Log.d(TAG, it) }
                            }
                        }

                        override fun onResult(hypothesis: String?) {
                            hypothesis?.let {
                                prepareWords(it,"onResult")
                            }
                        }
                        override fun onFinalResult(hypothesis: String?) {
                            hypothesis?.let {
                                prepareWords(it,"onFinalResult")
                            }
                            cont.resume(recMapRes)
                            speechStreamService.stop()
                        }

                        override fun onError(exception: java.lang.Exception?) {
                            //cont.resume(recMapRes)
                            Log.d(TAG, "Error: ${exception?.printStackTrace()}")
                        }

                        override fun onTimeout() {
                            cont.resume(recMapRes)
                        }

                    })
                }
            } catch (e: CancellationException) {
                // Handle cancellation
                println("Coroutine cancelled")
            } finally {

            }
        }

    }

    fun getCurrentState(): Int {
        return currentState
    }


}